"use client";
import {
  ReactFlow,
  Node,
  Edge,
  useNodesState,
  useEdgesState,
  useReactFlow,
  Background,
  NodeTypes,
} from "@xyflow/react";
import { useEffect, useMemo } from "react";
import ImageNode from "./image-node";
import ModelNode from "./model-node";

interface OutLoopFlowProps {
  inputImageUrl: string;
  outputImageUrl: string;
  chosenModels: string[]; // 检测、跟踪、匹配
}

export default function OutLoopFlow({
  inputImageUrl,
  outputImageUrl,
  chosenModels,
}: OutLoopFlowProps) {
  const [nodes, setNodes, onNodesChange] = useNodesState<Node>([]);
  const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([]);
  const { fitView } = useReactFlow();

  // 定义节点类型
  const nodeTypes: NodeTypes = useMemo(
    () => ({
      imageNode: ImageNode,
      modelNode: ModelNode,
    }),
    []
  );

  useEffect(() => {
    const newNodes: Node[] = [];
    const newEdges: Edge[] = [];

    // 输入图片节点
    newNodes.push({
      id: "input-images-node",
      type: "imageNode",
      position: { x: 0, y: 0 },
      data: { imageUrl: inputImageUrl, title: "输入图片" },
    });

    newNodes.push({
      id: "detection-model-node",
      type: "modelNode",
      position: { x: 250, y: 0 },
      data: {
        taskName: "检测模型",
        chosenModel: chosenModels[0],
      },
    });

    newNodes.push({
      id: "tracking-model-node",
      type: "modelNode",
      position: { x: 600, y: 0 },
      data: {
        taskName: "跟踪模型",
        chosenModel: chosenModels[1],
      },
    });

    newNodes.push({
      id: "search-model-node",
      type: "modelNode",
      position: { x: 950, y: 0 },
      data: {
        taskName: "匹配模型",
        chosenModel: chosenModels[2],
      },
    });

    // 输出图片节点
    newNodes.push({
      id: "output-image-node",
      type: "imageNode",
      position: { x: 1300, y: 0 }, // 纵向排列，y轴增加
      data: { imageUrl: outputImageUrl, title: "输出图片" },
    });

    // 创建边
    newEdges.push({
      id: "e1-2",
      source: "input-images-node",
      target: "detection-model-node",
      animated: true,
    });
    newEdges.push({
      id: "e2-3",
      source: "detection-model-node",
      target: "tracking-model-node",
      animated: true,
    });
    newEdges.push({
      id: "e3-4",
      source: "tracking-model-node",
      target: "search-model-node",
      animated: true,
    });
    newEdges.push({
      id: "e4-5",
      source: "search-model-node",
      target: "output-image-node",
      animated: true,
    });

    setNodes(newNodes);
    setEdges(newEdges);
  }, [inputImageUrl, outputImageUrl, chosenModels, setNodes, setEdges]);

  // 节点和边变化后自动 fitView
  useEffect(() => {
    fitView({ padding: 0.2 });
  }, [nodes, edges, fitView]);

  return (
    <div className="w-full h-full rounded-md border">
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        nodeTypes={nodeTypes}
        fitView
      >
        <Background />
      </ReactFlow>
    </div>
  );
}
